
/*
    Copyright(C) 2010 Ying-Chun Liu(PaulLiu). All rights reserved.

    This program is free software; you can redistribute it and/or modify
    it under the terms of the GNU General Public License as published by
    the Free Software Foundation; either version 2 of the License, or
    (at your option) any later version.

    This program is distributed in the hope that it will be useful,
    but WITHOUT ANY WARRANTY; without even the implied warranty of
    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
    GNU General Public License for more details.

    You should have received a copy of the GNU General Public License
    along with this program; if not, write to the Free Software
    Foundation, Inc., 51 Franklin St, Fifth Floor, Boston, MA  02110-1301  USA
*/
package org.debian.paulliu.net;

import java.io.*;
import java.util.*;
import java.net.*;
import java.util.logging.*;

public class StunClient {
    private InetSocketAddress serverSocketAddress;
    private String myPublicIP=null;
    private int myPublicPort=0;
    private String myLocalIP=null;
    private int myLocalPort=0;
    DatagramSocket socket=null;

    public enum MessageType {
	BINDING_REQUEST,
	    BINDING_RESPONSE,
	    BINDING_ERROR_RESPONSE,
	    SHARED_SECRET_REQUEST,
	    SHARED_SECRET_RESPONSE,
	    SHARED_SECRET_ERROR_RESPONSE
	    }

    /**
     * Class constructor
     * Assign the STUN server ip/port for this class
     *
     * @param serverSocketAddress STUN server ip/port
     */
    public StunClient(InetSocketAddress serverSocketAddress) {
	this.serverSocketAddress = serverSocketAddress;
    }

    private DatagramPacket GenerateSTUNPackage(MessageType type, Hashtable<String, java.lang.Object> attributes) {
	byte[] msgbuf;
	int msgbuflen = 20;

	if (attributes != null) {
	    for (Enumeration<String> i=attributes.keys(); i.hasMoreElements(); ) {
		String key;
		key = i.nextElement();

		if (key=="MAPPED-ADDRESS") {
		    msgbuflen += 12;
		} else if (key=="RESPONSE-ADDRESS") {
		    msgbuflen += 12;
		} else if (key=="CHANGED-ADDRESS") {
		    msgbuflen += 12;
		} else if (key=="CHANGE-REQUEST") {
		    msgbuflen += 8;
		} else if (key=="USERNAME") {
		    String username;
		    int username_len;
		    username = (String)attributes.get(key);
		    username_len = username.length();
		    username_len = ((username_len+3)/4)*4;
		    msgbuflen += username_len+4;
		} else if (key=="PASSWORD") {
		    String password;
		    int password_len;
		    password = (String)attributes.get(key);
		    password_len = password.length();
		    password_len = ((password_len+3)/4)*4;
		    msgbuflen += password_len+4;
		} else if (key=="MESSAGE-INTEGRITY") {
		    msgbuflen += 24;
		} else if (key=="ERROR-CODE") {
		    
		}
	    }
	}
	return null;
	
    }

    /**
     * Execute STUN. Prepare Public IP/Port and a socket.
     *
     */
    public void perform() {
	InetAddress serverAddr=null;
	DatagramPacket p;
	DatagramPacket d;
	byte[] buf;
	byte[] transactionID = new byte[16];
	byte[] myIPByte=null;
	int i,j;
	boolean checkSourceAddressFlag=false;
	myIPByte = new byte[4];
	serverAddr = this.serverSocketAddress.getAddress();
	if (serverAddr == null) {
	    Logger.getLogger("org.debian.paulliu").log(Level.SEVERE,"Unknown host "+this.serverSocketAddress.getHostName());
	}

	try {
	    socket = new DatagramSocket();
	} catch (java.net.SocketException e) {
	    Logger.getLogger("org.debian.paulliu").log(Level.SEVERE,"STUN: Error on DatagramSocket creation");
	}
	socket.connect(serverAddr, this.serverSocketAddress.getPort());
	this.myLocalIP = this.socket.getLocalAddress().getHostAddress();
	this.myLocalPort = this.socket.getLocalPort();
	socket.disconnect();
	socket.close();
	socket=null;

	try {
	    socket = new DatagramSocket(this.myLocalPort);
	} catch (java.net.SocketException e) {
	    Logger.getLogger("org.debian.paulliu").log(Level.SEVERE,"STUN: Error on DatagramSocket creation 2nd time");
	}
	try {
	    socket.setReuseAddress(true);
	} catch (java.net.SocketException e) {
	    Logger.getLogger("org.debian.paulliu").log(Level.SEVERE,"STUN: Error on setting SO_REUSEADDR");
	}
	try {
	    InetSocketAddress bindaddr = new InetSocketAddress(this.myLocalPort);
	    if (!socket.isBound()) {
		socket.bind(bindaddr);
	    }
	} catch (java.net.SocketException e) {
	    Logger.getLogger("org.debian.paulliu").log(Level.SEVERE,"STUN: Error on DatagramSocket bind on port: "+Integer.toString(this.myLocalPort));
	}

	buf = new byte[20];

	buf[0] = 0x00;
	buf[1] = 0x01;

	for (i=0; i<transactionID.length; i++) {
	    transactionID[i] = ((byte)(Math.random()*256));
	}
	for (i=4; i<20; i++) {
	    buf[i] = (transactionID[i-4]);
	}
	p = new DatagramPacket(buf,buf.length, serverAddr, this.serverSocketAddress.getPort());

	try {
	    socket.send(p);
	} catch (java.io.IOException e) {
	    Logger.getLogger("org.debian.paulliu").log(Level.SEVERE,"STUN: Send error");
	    return;
	}

	d = new DatagramPacket(new byte[512],512);

	try {
	    socket.receive(d);
	} catch (java.io.IOException e) {
	    Logger.getLogger("org.debian.paulliu").log(Level.SEVERE,"STUN: Receive error");
	    return;
	}

	for (i=0; i<d.getData().length; ) {
	    if (i==0 && i+20-1 < d.getData().length) {
		/* header */
		boolean checkIDFlag=true;
		if (d.getData()[i+0] != 1 || d.getData()[i+1] != 1) {
		    Logger.getLogger("org.debian.paulliu").log(Level.SEVERE,"STUN: Receive not reply");
		    myIPByte=null;
		    break;
		}
		for (j=0; j<transactionID.length; j++) {
		    if (d.getData()[i+4+j] != transactionID[j]) {
			checkIDFlag=false;
			break;
		    }
		}
		if (!checkIDFlag) {
		    Logger.getLogger("org.debian.paulliu").log(Level.SEVERE,"STUN: Transaction ID not match");
		    myIPByte=null;
		    break;
		}
		i+=20;
	    } else if (i+4-1 < d.getData().length) {
		if (d.getData()[i] == 0x00 && d.getData()[i+1] == 0x01) {
		    /* mapped address */
		    int length;
		    length = ((int)(d.getData()[i+2]&0xFF)) *256+ ((int)(d.getData()[i+3] & 0xFF));
		    for (j=0; j<4; j++) {
			myIPByte[j] = d.getData()[i+8+j];
		    }
		    this.myPublicPort = ((int)(d.getData()[i+6]&0xFF)) *256+ ((int)(d.getData()[i+7] & 0xFF));
		    i+=length+4;
		} else if (d.getData()[i] == 0x00 && d.getData()[i+1] == 0x04) {
		    /* source address */
		    int length;
		    byte[] serverSourceAddress = new byte[4];
		    InetAddress serverSourceAddr = null;
		    length = ((int)(d.getData()[i+2]&0xFF)) *256+ ((int)(d.getData()[i+3] & 0xFF)) ;
		    for (j=0; j<serverSourceAddress.length; j++) {
			serverSourceAddress[j] = d.getData()[i+8+j];
		    }
		    try {
			serverSourceAddr = InetAddress.getByAddress(serverSourceAddress);
		    } catch (java.net.UnknownHostException e) {
			Logger.getLogger("org.debian.paulliu").log(Level.SEVERE,"Cannot check server SOURCE-ADDRESS");
			myIPByte = null;
			break;
		    }
		    if (!serverSourceAddr.equals(serverAddr)) {
			Logger.getLogger("org.debian.paulliu").log(Level.SEVERE,"Attacker: SOURCE-ADDRESS modified");
			myIPByte = null;
			break;
		    }
		    checkSourceAddressFlag=true;
		    i+=length+4;
		} else {
		    int length;
		    length = ((int)(d.getData()[i+2]&0xFF)) *256+ ((int)(d.getData()[i+3] & 0xFF)) ;
		    i+=length+4;
		}
	    }
	}
	if (!checkSourceAddressFlag) {
	    Logger.getLogger("org.debian.paulliu").log(Level.SEVERE,"No SOURCE-ADDRESS");
	    myIPByte=null;
	}
	if (myIPByte != null) {
	    myPublicIP = Integer.toString(((int)(myIPByte[0]&0xFF)))+"."+Integer.toString(((int)(myIPByte[1]&0xFF)))+"."+Integer.toString(((int)(myIPByte[2]&0xFF)))+"."+Integer.toString(((int)(myIPByte[3]&0xFF)));
	}
    }

    /**
     * Get the socket
     */
    public DatagramSocket getSocket() {
	return this.socket;
    }

    /**
     * Get my public ip address from STUN server
     *
     * @return my ip address. String format. If anything is wrong, return null
     */
    public String getMyPublicIP() {
	return this.myPublicIP;
    }

    /**
     * Get my public port from STUN server
     *
     * @return my public port. 0 if something is wrong.
     */
    public int getMyPublicPort() {
	return this.myPublicPort;
    }


    /**
     * Get my private ip address
     *
     * @return my private ip address. String format.
     */
    public String getMyPrivateIP() {
	return this.myLocalIP;
    }

    /**
     * Get my private port from STUN server
     *
     * @return my private port. 0 if something is wrong.
     */
    public int getMyPrivatePort() {
	return this.myLocalPort;
    }


}

